classdef PoseEstimator < handle

	properties
	imagesT;
	imagesE;
	motions;
	featCalculator;
	featuresT;
	featuresTT;  % featuresT transposed
	coordinatesT;
	param;
	matcher;  % for kNN
	end

	methods
	function obj = PoseEstimator(varargin)
		root = CONFIG.SILHOUETTE_PATH;
		imc = ImageFileCollection('root', root);
		n = imc.nImages;

		% sets the seed of random number generator, for debugging purpose:
		rndgen('default');

		belongToTraining(n) = false;  % init
		belongToTraining(randsample(n, n / 2)) = true;

		% The "T" postfix here indicates the Training version, 
		% "E" indicates the Evaluation version.
		obj.imagesT = ImageFileCollection('root', root, 'database', ...
			imc.database(belongToTraining));
		obj.imagesE = ImageFileCollection('root', root, 'database', ...
			imc.database(~belongToTraining));

		% MotionCollection object can be shared
		obj.motions = MotionCollection(CONFIG.MOTION_PATH, 24, ...
			'partitionStr', 'exp');
		obj.imagesT.associatedMotions = obj.motions;
		obj.imagesE.associatedMotions = obj.motions;


		%--------------------------------%
		% for sparse representation
		addpath(CONFIG.LASSO_PATH);
		
		% equivalent to set gamma to 25 * 2 = 50
		obj.param.lambda = 10;  % TODO

		% the same as Elastic Net with the L2 penalty set to zero
		obj.param.lambda2 = 0;

		% non-negativity constraints
		obj.param.pos = true;

		% uses all the available CPUs/cores
		obj.param.numThreads = -1;

		% To minimize:
		%   1/2 |x - Dw|_2^2 + lambda |w|_1
		obj.param.mode = 2;
	end

	function [] = train(obj)
		obj.featCalculator = FeatureCalculator(obj.imagesT);
		
		% indicator = featureSelection(obj.featCalculator, 300);
		mask = FeatureCalculator.makeFeatureMask();
		indicator = mask.OM | mask.HU | mask.SC;
		fprintf('nDimSelected: %d\n', sum(indicator));
		obj.featCalculator.setFeatureSelector(indicator);

		obj.featuresT = obj.featCalculator.calculate('useSelector', true);
		obj.featuresTT = [];  % resets it

		% normalizes features, whose values lie within different dynamic ranges
		[obj.featuresT, obj.featCalculator.normalize] = ...
			FeatureCalculator.normalizerLerp(obj.featuresT);

		obj.coordinatesT = obj.featCalculator.calculatePoses();
		obj.matcher = ExhaustiveSearcher(obj.featuresT);
	end

	function averageError = evalKnn(obj, showImages)
		if ~exist('showImages', 'var')
			showImages = false;
		end

		% preserves other properties
		obj.featCalculator.imageCollection = obj.imagesE;
		
		averageError = 0;

		% to reduce memory occupation
		nImages = obj.imagesE.nImages;
		for fromTo = iterate(nImages, 150)
			idx = fromTo(1):fromTo(2);
			features = obj.featCalculator.calculate('useSelector', true, ...
				'indicesToCalc', idx, 'normalization', true);
			groundTruthCoord = obj.featCalculator.calculatePoses( ...
				'indicesToCalc', idx);
			% indices in the Training image collection
			nearestIdx = knnsearch(obj.matcher, features, 'K', 1);
			estimatedCoord = obj.coordinatesT(nearestIdx, :);

			if showImages
				subplot(1, 2, 1);
				imshow(obj.imagesE.at(idx(1)));
				title({'Query image:', obj.imagesE.database{idx(1)}}, ...
					'Interpreter', 'none');

				subplot(1, 2, 2);
				imshow(obj.imagesT.at(nearestIdx(1)));
				title({'Matched:', obj.imagesT.database{nearestIdx(1)}}, ...
					'Interpreter', 'none');
				fprintf('Query #%d in imagesE, get #%d in imagesT\n', ...
					idx(1), nearestIdx(1));
				fprintf('err = %f\n', poseDistance( ...
					estimatedCoord(1, :), groundTruthCoord(1, :)));
				1;%pause;
			end

			e = sum(poseDistance(estimatedCoord, groundTruthCoord));
			averageError = averageError + e;
		end
		
		averageError = averageError / nImages;
	end

	% - - - - - - - - - - - - - - - - - - - - - - - - %

	function averageError = evalKnnDP(obj, showImages)
		if ~exist('showImages', 'var')
			showImages = false;
		end

		% preserves other properties
		obj.featCalculator.imageCollection = obj.imagesE;
		
		%--------8<----------------
		function cost = transitionCost(p, q)
			sigma2 = 25;  % parameter sigma = 5
			cost = poseDistance(obj.coordinatesT(p, :), ...
				obj.coordinatesT(q, :));
			cost = exp(-cost / sigma2);
		end
		%---------------->8--------

		averageError = 0;

		% to reduce memory occupation
		nImages = obj.imagesE.nImages;
		optimizer = SequentialOptimizer(nImages, @transitionCost);

		for fromTo = iterate(nImages, 150)
			idx = fromTo(1):fromTo(2);
			features = obj.featCalculator.calculate('useSelector', true, ...
				'indicesToCalc', idx, 'normalization', true);
			groundTruthCoord = obj.featCalculator.calculatePoses( ...
				'indicesToCalc', idx);
			% indices in the Training image collection
			[nearestIdx, distances] = knnsearch(obj.matcher, features, 'K', 5);
			estimatedCoord = obj.coordinatesT(nearestIdx, :);

			for i = 1:length(idx)  % for each frame
				distances = distances / sum(distances);  % normalization
				optimizer.newFrame(distances, nearestIdx);

				estimatedCoords = obj.coordinatesT(indices, :);
				d = poseDistance(estimatedCoords, groundTruthCoord(idx(i), :));
				averageError = averageError + mean(weights * d);

			end

			if showImages
				subplot(1, 2, 1);
				imshow(obj.imagesE.at(idx(1)));
				title({'Query image:', obj.imagesE.database{idx(1)}}, ...
					'Interpreter', 'none');

				subplot(1, 2, 2);
				imshow(obj.imagesT.at(nearestIdx(1)));
				title({'Matched:', obj.imagesT.database{nearestIdx(1)}}, ...
					'Interpreter', 'none');
				fprintf('Query #%d in imagesE, get #%d in imagesT\n', ...
					idx(1), nearestIdx(1));
				fprintf('err = %f\n', poseDistance( ...
					estimatedCoord(1, :), groundTruthCoord(1, :)));
				1;%pause;
			end

			e = sum(poseDistance(estimatedCoord, groundTruthCoord));
			averageError = averageError + e;
		end

		choices = optimizer.makeChoices();
		
		averageError = averageError / nImages;
	end

	% - - - - - - - - - - - - - - - - - - - - - - - - %

	function averageError = evalSparse(obj, showImages)
		if ~exist('showImages', 'var')
			showImages = false;
		end

		% preserves other properties
		obj.featCalculator.imageCollection = obj.imagesE;
		
		if isempty(obj.featuresTT)
			obj.featuresTT = obj.featuresT';
		end

		%--------8<----------------
		function cost = transitionCost(p, q)
			sigma2 = 25;  % parameter sigma = 5
			cost = poseDistance(obj.coordinatesT(p, :), ...
				obj.coordinatesT(q, :));
			cost = exp(-cost / sigma2);
		end
		%---------------->8--------

		averageError = 0;
	
		nImages = obj.imagesE.nImages;
		% optimizer = SequentialOptimizer(nImages, transitionCost);

		% to reduce memory occupation
		for fromTo = iterate(nImages, 150)
			idx = fromTo(1):fromTo(2);
			features = obj.featCalculator.calculate('useSelector', true, ...
				'indicesToCalc', idx, 'normalization', true);
			groundTruthCoord = obj.featCalculator.calculatePoses( ...
				'indicesToCalc', idx);

			% solves the following optimization problem:
			% w = arg min(|x - D . w|_2 + gamma * |w|_1);
			w = mexLasso(features', obj.featuresTT, obj.param);
			% note that `w` is sparse, `ind` is full
			[w, ind] = sort(w, 'descend');
			nCandidates = full(sum(w ~= 0));  % 1-by-n

			for i = 1:size(nCandidates, 2)  % for each frame
				mask = 1:nCandidates(i);
				% fprintf('%d %d\n', fromTo(1), i);

				weights = full(w(mask, i)');
				weights = weights / sum(weights);  % normalization
				indices = ind(mask, i)';
				% optimizer.newFrame(weights, indices);

				estimatedCoords = obj.coordinatesT(indices, :);
				d = poseDistance(estimatedCoords, groundTruthCoord(i, :));
				averageError = averageError + mean(weights * d);

			end

			nearestIdx = indices(1);  % 1-by-1 here
			if showImages
				subplot(1, 2, 1);
				imshow(obj.imagesE.at(idx(1)));
				title({'Query image:', obj.imagesE.database{idx(1)}}, ...
					'Interpreter', 'none');

				subplot(1, 2, 2);
				imshow(obj.imagesT.at(nearestIdx(1)));
				title({'Matched:', obj.imagesT.database{nearestIdx(1)}}, ...
					'Interpreter', 'none');
				fprintf('Query #%d in imagesE, get #%d in imagesT\n', ...
					idx(1), nearestIdx(1));
				fprintf('err = %f\n', poseDistance( ...
					estimatedCoords(1, :), groundTruthCoord(1, :)));
				1;%pause;
			end

			% e = sum(poseDistance(estimatedCoords, groundTruthCoord));
			% averageError = averageError + e;
		end
		
		% choices = optimizer.makeChoices();

		averageError = averageError / nImages;
	end

	% - - - - - - - - - - - - - - - - - - - - - - - - %

	function averageError = evalRandom(obj)
		rndgen('shuffle');

		% preserves other properties
		obj.featCalculator.imageCollection = obj.imagesE;

		nImages = obj.imagesE.nImages;
		nearestIdx = randi(nImages, nImages, 1);
		groundTruthCoord = obj.featCalculator.calculatePoses();
		estimatedCoord = obj.coordinatesT(nearestIdx, :);
		averageError = ...
			sum(poseDistance(estimatedCoord, groundTruthCoord)) / nImages;

	end

	function [] = renderKnn(obj)
		% preserves other properties
		obj.featCalculator.imageCollection = obj.imagesE;

		% turns off figure displaying temporarily
		set(gcf, 'Visible', 'off');

		% to reduce memory occupation
		nImages = obj.imagesE.nImages;
		for fromTo = iterate(nImages, 150)
			idx = fromTo(1):fromTo(2);
			features = obj.featCalculator.calculate('useSelector', true, ...
				'indicesToCalc', idx, 'normalization', true);
			% indices in the Training image collection
			nearestIdx = knnsearch(obj.matcher, features, 'K', 1);
			estimatedCoord = obj.coordinatesT(nearestIdx, :);


			for i = 1:length(idx)
				destSuffix = strrep(obj.imagesE.database{idx(i)}, ...
					'.jpg', '.png');
				destPath = ['F:\poseVideo\videoOut\', destSuffix];
				destDir = fileparts(destPath);

				% ingores existing files
				if exist(destPath, 'file') == 2
					continue;
				end

				coord = estimatedCoord(i, :);
				SkeletonDrawer.draw(coord);

				if exist(destDir, 'file') ~= 7
					mkdir(destDir);
				end
				saveas(gcf, destPath);
			end
		end
		set(gcf, 'Visible', 'on');
	end

	function averageError = evalBest(obj)
		% preserves other properties
		obj.featCalculator.imageCollection = obj.imagesE;
		averageError = 0;

		nImages = obj.imagesE.nImages;
		for i = 1:nImages
			[~, e] = obj.exhaustiveMatch(i);
			averageError = averageError + e;
		end

		averageError = averageError / nImages;
	end

	function [matchedId, distance] = exhaustiveMatch(obj, queryId)
		p = obj.imagesE.associatedPoseAt(queryId);
		dis = poseDistance(p, obj.coordinatesT);
		[distance, matchedId] = min(dis);
	end
	end

	methods (Static)
	function [] = saveToDisk(pe)
		assert(isa(pe, 'PoseEstimator'));
		pe.featuresTT = [];  % less space
		save('poseEstimatorCached.mat', '-v6', 'pe');
	end

	function [] = loadFromDisk()
		load('poseEstimatorCached.mat');
	end
	end
end

function a = iterate(n, k)
% iterate(32, 7) = 
% 	1     8    15    22    29
% 	7    14    21    28    32

	temp = (1:ceil(n/k)) * k;
	a = [temp - (k - 1); temp];
	a(end) = n;
end

function [] = rndgen(option)
	switch option
	case 'default'
		% MATLAB's start-up setting
		seed = 5489;
	case 'shuffle'
		seed = sum(100 * clock);
	otherwise
		seed = option;
	end

	stream = RandStream('mt19937ar', 'Seed', seed);
	RandStream.setDefaultStream(stream);
end
